package com.shao.loadBalance.impl;

import com.shao.SrpcBootstrap;
import com.shao.loadBalance.AbstractLoadBalancer;
import com.shao.loadBalance.Selector;
import com.shao.transport.message.SrpcRequest;
import lombok.Data;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.security.MessageDigest;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * @author PaoFu
 * @version 1.0
 * @data 2025/5/6 14:48
 */
@Slf4j
public class ConsistentHashLoadBalancer extends AbstractLoadBalancer {

    @Override
    protected Selector getSelector(List<InetSocketAddress> addresses) {
        return new ConsistentHashSelector(addresses, 256);
    }

    @Data
    @Accessors(chain = true)
    private static class ConsistentHashSelector implements Selector {

        private TreeMap<Integer, InetSocketAddress> circle = new TreeMap<>();

        private List<InetSocketAddress> socketAddressesList;

        private int virtualNodes;

        public ConsistentHashSelector(List<InetSocketAddress> socketAddressesList, int virtualNodes) {
            // 需要把节点转化成虚拟节点
            this.virtualNodes = virtualNodes;
            for (InetSocketAddress inetSocketAddress : socketAddressesList) {
                addNode2Circle(inetSocketAddress);
            }
        }

        @Override
        public InetSocketAddress getNext() {
            SrpcRequest srpcRequest = SrpcBootstrap.REQUEST_THREAD_LOCAL.get();
            String requestId = Long.toString(srpcRequest.getRequestId());
            int hash = hash(requestId);
            if (!circle.containsKey(hash)) {
                // 如果哈希环不包含当前请求的哈希值则查找最近的一个
                SortedMap<Integer, InetSocketAddress> tailMap = circle.tailMap(hash);
                hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
            }
            log.debug("使用哈希负载均衡获取下一个服务节点【{}】", circle.get(hash));
            return circle.get(hash);
        }

        /**
         * 将每个节点挂载到哈希环上
         * @param inetSocketAddress 节点的地址
         */
        private void addNode2Circle(InetSocketAddress inetSocketAddress) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(inetSocketAddress + "-" + i);
                circle.put(hash, inetSocketAddress);
                if (log.isDebugEnabled()) {
                    log.debug("hash为【{}】的节点挂载到了哈希环上", hash);
                }
            }
        }

        /**
         * 将指定节点在哈希环上移除
         * @param inetSocketAddress 节点的地址
         */
        private void removeNode2Circle(InetSocketAddress inetSocketAddress) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(inetSocketAddress + "-" + i);
                circle.remove(hash, inetSocketAddress);
            }
        }

        private int hash(String msg) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");
            } catch (Exception e) {
                log.error("创建MD5实例失败");
                throw new RuntimeException(e);
            }
            byte[] md5Hash = md.digest(msg.getBytes());
            byte[] first4Bytes = getFirst4Bytes(md5Hash);
            return bytesToIntBE(first4Bytes);
        }

        /**
         * 提取前4个字节
         * @param md5Hash md5哈希结果
         * @return 返回前4个字节
         */
        private byte[] getFirst4Bytes(byte[] md5Hash) {
            byte[] first4 = new byte[4];
            System.arraycopy(md5Hash, 0, first4, 0, 4);
            return first4;
        }
        /**
         * 大端模式字节转int
         */
        public static int bytesToIntBE(byte[] bytes) {
            if (bytes.length < 4) {
                throw new IllegalArgumentException("字节数组需至少4个元素");
            }
            return ((bytes[0] & 0xFF) << 24) |  // 消除符号影响，左移24位
                    ((bytes[1] & 0xFF) << 16) |
                    ((bytes[2] & 0xFF) << 8)  |
                    (bytes[3] & 0xFF);
        }
    }
}
